package com.shardingsphere.example.aspect;


import com.shardingsphere.example.annotation.ShardingTable;
import com.shardingsphere.example.sharding.ShardingContextHolder;
import com.shardingsphere.example.sharding.ModHashShardingStrategy;
import com.shardingsphere.example.sharding.RangeShardingStrategy;
import com.shardingsphere.example.sharding.ShardingStrategy;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 分表服务切面
 */
@Aspect
@Component
public class ShardingAspect {

    private final Map<String, ShardingStrategy> strategyMap = new ConcurrentHashMap<>();

    public ShardingAspect() {
        strategyMap.put("MOD_HASH", new ModHashShardingStrategy());
        strategyMap.put("RANGE", new RangeShardingStrategy());
    }

    @Around("@annotation(shardingTable)")
    public Object around(ProceedingJoinPoint joinPoint, ShardingTable shardingTable) throws Throwable {
        try {
            // 解析参数，获取分片键值
            Object shardingValue = getShardingValue(joinPoint, shardingTable);

            // 计算实际表名
            String strategyType = shardingTable.strategy().name();
            ShardingStrategy strategy = strategyMap.get(strategyType);
            if (strategy == null) {
                strategy = strategyMap.get("MOD_HASH");
            }

            String actualTableName = strategy.getActualTableName(
                    shardingTable.logicTable(),
                    shardingValue,
                    shardingTable.tableCount());

            // 设置到上下文
            ShardingContextHolder.setTableName(actualTableName);

            // 执行目标方法
            return joinPoint.proceed();
        } finally {
            // 清除上下文
            ShardingContextHolder.clear();
        }
    }

    /**
     * 获取分片键值
     */
    private Object getShardingValue(ProceedingJoinPoint joinPoint, ShardingTable shardingTable) {
        String shardingKey = shardingTable.shardingKey();
        if (shardingKey == null || shardingKey.isEmpty()) {
            return null;
        }

        // 获取方法参数
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        String[] parameterNames = signature.getParameterNames();
        Object[] args = joinPoint.getArgs();

        // 查找对应的参数
        for (int i = 0; i < parameterNames.length; i++) {
            if (shardingKey.equals(parameterNames[i])) {
                return args[i];
            }
        }

        // 如果通过参数名没找到，尝试通过对象属性查找
        for (Object arg : args) {
            if (arg != null) {
                try {
                    java.lang.reflect.Field field = arg.getClass().getDeclaredField(shardingKey);
                    field.setAccessible(true);
                    return field.get(arg);
                } catch (Exception e) {
                    // 忽略异常，继续查找
                }
            }
        }

        return null;
    }
}

